
"""
角色抽取
"""
from models import Novel, Chapter
from pydantic import Field
from core.utils import BaseAnalyzer, prompt_dir
from core.alone_analyze.simply_extract import SimplyExtractor
from typing import List, Any


class CharactersExtractor(BaseAnalyzer):

    def __init__(self):
        super().__init__(prompt_file=prompt_dir / 'context/c2. 人物信息抽取.md')
        self.simply_extractor = SimplyExtractor()

    def invoke(
            self,
            novel: Novel = Field(description="小说"),
            archived_characters: List = Field(description="人物的历史信息"),
            chapter: Chapter = Field(description="章节信息")
    ):
        # 1. 结合历史信息和当前章节内容对人物塑造进行分析
        prompt = self.template_prompt.format(
            inputs={
                "chapter_content": chapter.chapter_content,
                "archived_characters": '\n'.join([str(character) for character in archived_characters]),
                "character_profile": chapter.character_profile
            },
            remove_template_variables=True
        )
        chapter.characters = self.analyze(prompt)

        # 2. 更新人物信息到到全局记录
        self.update_characters(novel, chapter.characters)
        return chapter.characters

    @staticmethod
    def update_characters(novel: Novel, characters: Any) -> None:
        """ 更新小说的人物信息 """
        if not isinstance(characters, list):
            return
        for character in characters:
            if not isinstance(character, dict):
                continue
            name = character.get("名称", None)
            if name is None:
                continue
            novel.characters[name] = character
            old_name = character.get("旧名称", None)
            if old_name is None or old_name == "":
                continue
            elif old_name != name:
                novel.characters.pop(old_name, None)
